from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA
from langchain_openai import AzureChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain_openai.embeddings import AzureOpenAIEmbeddings  # 导入嵌入模型
from tool import get_azure_endpoint,get_api_version,get_api_key

if __name__ == '__main__':
    # Build prompt
    template = """使用以下上下文片段来回答最后的问题。如果你不知道答案，只需说不知道，不要试图编造 答案。答案最多使用三个句子。尽量简明扼要地回答。在回答的最后一定要说"感谢您的提问!"
    {context}
    问题:{question}
    有用的回答:"""
    QA_CHAIN_PROMPT = PromptTemplate.from_template(template)

    llm = AzureChatOpenAI(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        azure_deployment="gpt-4o-mini",  # 重命名为 azure_deployment
        openai_api_version=get_api_version(),  # 参数名不变
        openai_api_key=get_api_key(),
        openai_api_type="azure",
    )
    embedding = AzureOpenAIEmbeddings(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        model="text-embedding-3-small",  # 重命名为 azure_deployment
        api_key=get_api_key(),
        api_version=get_api_version()
    )
    vectordb_chinese = Chroma(
        persist_directory="./docs/chroma/matplotlib",
        embedding_function=embedding,
    )

    # Run chain
    qa_chain = RetrievalQA.from_chain_type(
        llm,
        retriever=vectordb_chinese.as_retriever(),
        return_source_documents=True,
        chain_type_kwargs={"prompt": QA_CHAIN_PROMPT}
    )

    question = "这门课需要学习 python 吗"
    result = qa_chain({"query": question})
    print(result["result"])
    print(result["source_documents"][0])